package com.by.log;

import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 分布式追踪上下文工具类
 * 用于在分布式系统中传递和管理追踪信息
 */
public class TraceContextUtil {
    private static final String TRACE_ID_KEY = "X-Trace-Id";
    private static final String SPAN_ID_KEY = "X-Span-Id";
    private static final String PARENT_SPAN_ID_KEY = "X-Parent-Span-Id";
    private static final String SERVICE_NAME_KEY = "X-Service-Name";
    private static final String METHOD_NAME_KEY = "X-Method-Name";
    
    private static final ThreadLocal<TraceData> TRACE_CONTEXT = new ThreadLocal<TraceData>() {
        @Override
        protected TraceData initialValue() {
            return new TraceData();
        }
    };

    /**
     * 追踪数据类
     */
    public static class TraceData {
        private String traceId;
        private String spanId;
        private String parentSpanId;
        private String serviceName;
        private String methodName;
        private final Map<String, String> attributes = new ConcurrentHashMap<>();
        
        public String getTraceId() {
            return traceId;
        }
        
        public void setTraceId(String traceId) {
            this.traceId = traceId;
        }
        
        public String getSpanId() {
            return spanId;
        }
        
        public void setSpanId(String spanId) {
            this.spanId = spanId;
        }
        
        public String getParentSpanId() {
            return parentSpanId;
        }
        
        public void setParentSpanId(String parentSpanId) {
            this.parentSpanId = parentSpanId;
        }
        
        public String getServiceName() {
            return serviceName;
        }
        
        public void setServiceName(String serviceName) {
            this.serviceName = serviceName;
        }
        
        public String getMethodName() {
            return methodName;
        }
        
        public void setMethodName(String methodName) {
            this.methodName = methodName;
        }
        
        public Map<String, String> getAttributes() {
            return attributes;
        }
        
        public String getAttribute(String key) {
            return attributes.get(key);
        }
        
        public void setAttribute(String key, String value) {
            attributes.put(key, value);
        }
        
        public void removeAttribute(String key) {
            attributes.remove(key);
        }
    }

    /**
     * 获取当前线程的追踪数据
     * @return 追踪数据
     */
    public static TraceData getCurrentTraceData() {
        return TRACE_CONTEXT.get();
    }

    /**
     * 清除当前线程的追踪上下文
     */
    public static void clear() {
        TRACE_CONTEXT.remove();
    }

    /**
     * 从请求头中提取追踪信息并设置到当前上下文
     * @param headers 请求头
     */
    public static void extractTraceContext(Map<String, String> headers) {
        TraceData traceData = getCurrentTraceData();
        traceData.setTraceId(headers.get(TRACE_ID_KEY));
        traceData.setSpanId(headers.get(SPAN_ID_KEY));
        traceData.setParentSpanId(headers.get(PARENT_SPAN_ID_KEY));
        traceData.setServiceName(headers.get(SERVICE_NAME_KEY));
        traceData.setMethodName(headers.get(METHOD_NAME_KEY));
    }

    /**
     * 将当前追踪信息注入到请求头中
     * @param headers 请求头
     */
    public static void injectTraceContext(Map<String, String> headers) {
        TraceData traceData = getCurrentTraceData();
        if (traceData.getTraceId() != null) {
            headers.put(TRACE_ID_KEY, traceData.getTraceId());
        }
        if (traceData.getSpanId() != null) {
            headers.put(SPAN_ID_KEY, traceData.getSpanId());
        }
        if (traceData.getParentSpanId() != null) {
            headers.put(PARENT_SPAN_ID_KEY, traceData.getParentSpanId());
        }
        if (traceData.getServiceName() != null) {
            headers.put(SERVICE_NAME_KEY, traceData.getServiceName());
        }
        if (traceData.getMethodName() != null) {
            headers.put(METHOD_NAME_KEY, traceData.getMethodName());
        }
    }

    /**
     * 开始一个新的追踪
     * @param serviceName 服务名称
     * @param methodName 方法名称
     */
    public static void startTrace(String serviceName, String methodName) {
        TraceData traceData = getCurrentTraceData();
        traceData.setTraceId(generateTraceId());
        traceData.setSpanId(generateSpanId());
        traceData.setParentSpanId(null);
        traceData.setServiceName(serviceName);
        traceData.setMethodName(methodName);
    }

    /**
     * 开始一个新的跨度
     * @param serviceName 服务名称
     * @param methodName 方法名称
     */
    public static void startSpan(String serviceName, String methodName) {
        TraceData traceData = getCurrentTraceData();
        String parentSpanId = traceData.getSpanId();
        traceData.setSpanId(generateSpanId());
        traceData.setParentSpanId(parentSpanId);
        traceData.setServiceName(serviceName);
        traceData.setMethodName(methodName);
    }

    /**
     * 生成追踪ID
     * @return 追踪ID
     */
    public static String generateTraceId() {
        return UUID.randomUUID().toString().replace("-", "");
    }

    /**
     * 生成跨度ID
     * @return 跨度ID
     */
    public static String generateSpanId() {
        return String.valueOf(System.nanoTime());
    }
}